{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import xml.etree.ElementTree as ET;\n",
    "from Tools.Parsers.KdTreeXMLParser import KdTree\n",
    "\n",
    "import re, os\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "from vispy import app\n",
    "from vispy import scene\n",
    "from vispy.visuals.transforms import STTransform, NullTransform, AffineTransform, BaseTransform\n",
    "from vispy.color.colormap import *\n",
    "\n",
    "from transforms3d import quaternions, axangles\n",
    "\n",
    "def generateGridBoxLines(minP = np.array([-1.0,-1.0,-1.0]), \n",
    "              maxP = np.array([1.0,1.0,1.0]), \n",
    "              rotMatrix = None,\n",
    "              subdivs=None, boundaryBox = True):\n",
    "    \n",
    "    centerPos = (maxP + minP)/2.0;\n",
    "    extent = (maxP - minP);\n",
    "    \n",
    "    if boundaryBox:\n",
    "        nLines = 12;\n",
    "    else:\n",
    "        nLinse = 0;\n",
    "    # for each subdiv a quad is (4 lines) is drawn\n",
    "    if subdivs:\n",
    "       nLines += (subdivs[0]+1)*(subdivs[1]+1) + (subdivs[1]+1)*(subdivs[2]+1) + (subdivs[2]+1)*(subdivs[0]+1) - 4*3;\n",
    "\n",
    "    lines = np.zeros((nLines*2,3), dtype=np.float32);\n",
    "    \n",
    "    # plot lines\n",
    "    if boundaryBox:\n",
    "        lines[0:24,0:3] = 0.5* np.array([ (-1, -1, -1),\n",
    "                                        (1, -1, -1),\n",
    "                                        #\n",
    "                                        (1, -1, -1),\n",
    "                                        (1, 1, -1),\n",
    "                                        #\n",
    "                                        (1, 1, -1),\n",
    "                                        (-1, 1, -1), \n",
    "                                        #\n",
    "                                        (-1, 1, -1),\n",
    "                                        (-1, -1, -1), # end of first quad\n",
    "                                        #\n",
    "                                        (-1, -1, 1),\n",
    "                                        (1, -1, 1),\n",
    "                                        #\n",
    "                                        (1, -1, 1),\n",
    "                                        (1, 1, 1),\n",
    "                                        #\n",
    "                                        (1, 1, 1),\n",
    "                                        (-1, 1, 1), \n",
    "                                        #\n",
    "                                        (-1, 1, 1),\n",
    "                                        (-1, -1, 1), # end of second quad\n",
    "                                        #\n",
    "                                        (-1, -1, -1),\n",
    "                                        (-1, -1, 1),\n",
    "                                        #\n",
    "                                        (1, -1, -1),\n",
    "                                        (1, -1, 1),\n",
    "                                        #\n",
    "                                        (-1, 1, -1),\n",
    "                                        (-1, 1, 1), \n",
    "                                        #\n",
    "                                        (1, 1, -1),\n",
    "                                        (1, 1, 1)]); # end of vertical lines\n",
    "    \n",
    "    \n",
    "    # add subdivs if necessary\n",
    "\n",
    "    if subdivs is not None:\n",
    "        counter = 12*2; # start after the cube lines\n",
    "        dxyz = np.array([1.0 / subdivs[0],1.0 / subdivs[1],1.0 / subdivs[2]]);\n",
    "        \n",
    "        # y\n",
    "        # |    +   +   +   \n",
    "        # | *  *   *   *   *\n",
    "        # | *  *   *   *   *\n",
    "        # |    +   +   +  \n",
    "        # ----------------> x\n",
    "        #  * come first (, then the +\n",
    "        \n",
    "        # make x lines inside (*)\n",
    "        for k in range(1,int(subdivs[1])):\n",
    "            for l in range(0,int(subdivs[2])+1):\n",
    "                lines[counter:counter+2,0:3] = np.array([ (-0.5, -0.5 + k*dxyz[1],  -0.5 + l*dxyz[2]),\n",
    "                                                          (0.5, -0.5 + k*dxyz[1],  -0.5 + l*dxyz[2])  ]);\n",
    "                counter+=2;\n",
    "        # make the x lines (+)\n",
    "        for k in (0,subdivs[1]):\n",
    "            for l in range(1,int(subdivs[2])):\n",
    "                lines[counter:counter+2,0:3] = np.array([ (-0.5, -0.5 + k*dxyz[1],  -0.5 + l*dxyz[2]),\n",
    "                                                          (0.5, -0.5 + k*dxyz[1],  -0.5 + l*dxyz[2])  ]);\n",
    "                counter+=2;\n",
    "        \n",
    "        \n",
    "        # make y lines inside (*)     \n",
    "        for k in range(1,int(subdivs[2])):\n",
    "            for l in range(0,int(subdivs[0])+1):\n",
    "                lines[counter:counter+2,0:3] = np.array([ (-0.5+ l*dxyz[0],  -0.5,  -0.5 + k*dxyz[2]),\n",
    "                                                          (-0.5+ l*dxyz[0],  0.5,  -0.5 + k*dxyz[2])  ]);\n",
    "                counter+=2;\n",
    "        # make the y lines (+)\n",
    "        for k in (0,subdivs[2]):\n",
    "            for l in range(1,int(subdivs[0])):\n",
    "                lines[counter:counter+2,0:3] = np.array([ (-0.5+ l*dxyz[0],  -0.5,  -0.5 + k*dxyz[2]),\n",
    "                                                          (-0.5+ l*dxyz[0],  0.5,  -0.5 + k*dxyz[2])  ]);\n",
    "                counter+=2;\n",
    "                \n",
    "        # make z lines inside\n",
    "        for k in range(1,int(subdivs[0])):\n",
    "            for l in range(0,int(subdivs[1])+1):\n",
    "                lines[counter:counter+2,0:3] = np.array([ (-0.5+k*dxyz[0],  -0.5+l*dxyz[1],  -0.5),\n",
    "                                                          (-0.5+k*dxyz[0],  -0.5+l*dxyz[1],  0.5)    ]); \n",
    "                counter+=2;\n",
    "        # make the z lines (+)\n",
    "        for k in (0,subdivs[0]):\n",
    "            for l in range(1,int(subdivs[1])):\n",
    "                lines[counter:counter+2,0:3] = np.array([ (-0.5+k*dxyz[0],  -0.5+l*dxyz[1],  -0.5),\n",
    "                                                          (-0.5+k*dxyz[0],  -0.5+l*dxyz[1],  0.5)    ]); \n",
    "                counter+=2;\n",
    "        \n",
    "\n",
    "    lines[...] *= extent;\n",
    "    if rotMatrix is not None:\n",
    "        lines = np.dot(lines,rotMatrix.T); # [p1.T; p2.T] * R.transpose\n",
    "    lines[...] += centerPos;    \n",
    "    \n",
    "    return lines;\n",
    "\n",
    "class PlotKdTree:\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "   \n",
    "    \n",
    "    def plot(self,kdTree,plotPoints,plotSubDivs):\n",
    "        \n",
    "        topoColor=[1,0,0]\n",
    "        opts={ 'topo':{'lc':topoColor,'mc':topoColor} , \n",
    "               'points':{'lc':'','ls':'','mc':topoColor,'ms':'.','mw':2}, \n",
    "               'subdivs' :{'ls':'+'} }\n",
    "        \n",
    "        canvas = scene.SceneCanvas(size=(800, 600), show=True, keys='interactive', )\n",
    "        viewBox = canvas.central_widget.add_view(); # create view box\n",
    "        viewBox.camera = scene.TurntableCamera(elevation=30, azimuth=30, up='+z')\n",
    "        axis = scene.visuals.XYZAxis(parent=viewBox.scene);       \n",
    "        \n",
    "        # Plotting KDTREE\n",
    "        levels = kdTree.leafs.keys();\n",
    "\n",
    "        cm = self.generateColorMap(max=max(levels)-min(levels));\n",
    "\n",
    "        # plot root box\n",
    "        lines = generateGridBoxLines(minP=kdTree.aabb.min, maxP=kdTree.aabb.max);\n",
    "        colors = np.ones((len(lines),4), dtype=np.float32); colors[...,0:3]*= opts['topo']['mc'];\n",
    "\n",
    "        if plotSubDivs:\n",
    "\n",
    "            # add all leafs\n",
    "            for level,leafs in kdTree.leafs.items():\n",
    "                for leaf in leafs:\n",
    "                    l = generateGridBoxLines(minP=leaf.aabb.min, maxP=leaf.aabb.max);\n",
    "                    color = np.random.rand(1,3);\n",
    "                    c = np.ones((len(l),4), dtype=np.float32); c[...,0:3]= color;#cm.to_rgba(level);\n",
    "                    colors = np.concatenate(( colors,  c));\n",
    "                    lines = np.concatenate(( lines,  l));\n",
    "                    #scene.visuals.Text(str(leaf.idx), parent=viewBox.scene, color=color, pos = leaf.aabb.center()[0:3]);\n",
    "#                 # add all lines for the aabbs\n",
    "#                 counter = 0;\n",
    "#                 for level, aabbs in kdTree.aabbTree.items():\n",
    "#                     if level == 0:\n",
    "#                         continue;\n",
    "#                     for aabb in aabbs:\n",
    "#                         l = generateGridBoxLines(minP=aabb['min'], maxP=aabb['max']);\n",
    "#                         c = np.ones((len(l),4), dtype=np.float32); c*= cm.to_rgba(level);\n",
    "#                         colors = np.concatenate(( colors,  c));\n",
    "#                         lines = np.concatenate(( lines,  l));\n",
    "#                         counter+=1;\n",
    "\n",
    "        scene.visuals.Line(pos=lines,color=colors,\n",
    "                               parent=viewBox.scene,\n",
    "                               name=\"kdTree\",\n",
    "                               connect='segments', method='gl', width=2)    \n",
    "\n",
    "        app.run();\n",
    "        \n",
    "        \n",
    "    def generateColorMap(self,max, min=0):\n",
    "        hsl = colorMap = get_colormap('hsl') \n",
    "        colors = hsl.map( np.linspace(0,1,max) )\n",
    "        return colors;\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "aligned: True\n",
      "Root AABB:  <Tools.Parsers.AABB.AABB object at 0x7f6dc27a4ef0>\n",
      "A_IK: \n",
      " [[ 1.  0.  0.]\n",
      " [ 0.  1.  0.]\n",
      " [ 0.  0.  1.]]\n"
     ]
    }
   ],
   "source": [
    "    \n",
    "\n",
    "file = \"KdTreeResults.xml\"\n",
    "tree = ET.parse(file)\n",
    "root = tree.getroot()\n",
    "\n",
    "kdTree = KdTree.parseFromXML(root)\n",
    "p = PlotKdTree()\n",
    "p.plot(kdTree,False,True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
